import cv2
import numpy as np
import matplotlib.pyplot as plt
from scipy.signal import find_peaks
from skimage.filters import gabor_kernel
from skimage.util import img_as_ubyte
import pytesseract
import re
import os


def detect_scale_bar(gray_img, scale_region_size=(40, 100), fallback_nm=10.0):
    h, w = gray_img.shape
    bar_h, bar_w = scale_region_size
    region = gray_img[h - bar_h:283, :bar_w]

    dark_mask = (region < 40).astype(np.uint8)
    contours, _ = cv2.findContours(dark_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    max_width = 0
    for cnt in contours:
        x, y, w, h = cv2.boundingRect(cnt)
        if w > max_width:
            max_width = w

    config = '--psm 6'
    ocr_text = pytesseract.image_to_string(region, config=config)
    print(f"🔍 OCR识别内容: '{ocr_text.strip()}'")

    match = re.search(r"(\d+(\.\d+)?)\s*(nm)?", ocr_text.lower())
    if match:
        scale_nm_raw = float(match.group(1))
        allowed = [5, 10, 20, 50, 100, 200]
        scale_nm = min(allowed, key=lambda x: abs(x - scale_nm_raw))
        print(f"✅ 提取数值: {scale_nm_raw} → 匹配最接近合法值: {scale_nm} nm")
    else:
        print(f"⚠️ OCR 未识别到有效比例尺文字，使用默认值 {fallback_nm} nm")
        scale_nm = fallback_nm

    if max_width == 0:
        raise ValueError("比例尺检测失败：未找到明显黑条。")

    print(f"📏 比例尺条像素长度: {max_width} px")
    nm_per_pixel = scale_nm / max_width
    print(f"📐 换算结果: {nm_per_pixel:.4f} nm/pixel")

    return nm_per_pixel


def apply_gabor_filter(gray_img, frequency=0.1):
    accum = np.zeros_like(gray_img, dtype=np.float32)
    for theta in np.linspace(0, np.pi, 8):
        kernel = np.real(gabor_kernel(frequency, theta=theta))
        filtered = cv2.filter2D(gray_img.astype(np.float32), -1, kernel)
        accum = np.maximum(accum, filtered)
    return img_as_ubyte(accum / accum.max())


def extract_profile_and_spacing(filtered_img, nm_per_pixel, num_profiles=5):
    h, w = filtered_img.shape
    spacing_list = []
    profiles = []
    peaks_list = []

    col_center = w // 2
    offsets = np.linspace(-10, 10, num_profiles).astype(int)

    for dx in offsets:
        col = col_center + dx
        profile = np.mean(filtered_img[:, col-1:col+1], axis=1)
        profile_inv = 255 - profile
        peaks, properties = find_peaks(profile_inv, distance=5, prominence=5)

        if len(peaks) >= 2:
            distances = np.diff(peaks)
            spacing_px = np.mean(distances)
            spacing_nm = spacing_px * nm_per_pixel
            spacing_list.append(spacing_nm)
        else:
            spacing_list.append(None)

        profiles.append(profile)
        peaks_list.append(peaks)

    valid_spacings = [s for s in spacing_list if s is not None]
    avg_spacing_nm = np.mean(valid_spacings) if valid_spacings else None

    return avg_spacing_nm, spacing_list, profiles, peaks_list


def find_best_frequency(gray_img, nm_per_pixel, num_profiles=5):
    best_score = -np.inf
    best_freq = None
    best_filtered = None

    for freq in np.linspace(0.05, 0.3, 10):
        filtered = apply_gabor_filter(gray_img, frequency=freq)
        avg_spacing, spacings, profiles, peaks_list = extract_profile_and_spacing(filtered, nm_per_pixel, num_profiles=num_profiles)

        valid_peaks = [len(p) for p in peaks_list if len(p) >= 3]
        if len(valid_peaks) == 0:
            continue

        peak_count_score = np.mean(valid_peaks)
        spacing_score = avg_spacing if avg_spacing else 0
        score = peak_count_score * spacing_score

        if score > best_score:
            best_score = score
            best_freq = freq
            best_filtered = filtered

    print(f"🔎 最佳Gabor频率: {best_freq:.3f}")
    return best_freq, best_filtered


def extract_spacing_fft(gray_img, nm_per_pixel):
    """
    使用 2D FFT 分析图像主频率，推测主条纹间距。

    参数：
        gray_img: 原始灰度图像（2D）
        nm_per_pixel: 每像素代表的实际长度

    返回：
        spacing_nm: 由主频率计算出的条纹间距（单位：nm）
    """
    f = np.fft.fft2(gray_img)
    fshift = np.fft.fftshift(f)
    magnitude_spectrum = np.abs(fshift)

    h, w = gray_img.shape
    cy, cx = h // 2, w // 2
    magnitude_spectrum[cy-5:cy+6, cx-5:cx+6] = 0  # 去除低频中心

    max_pos = np.unravel_index(np.argmax(magnitude_spectrum), magnitude_spectrum.shape)
    dy = max_pos[0] - cy
    dx = max_pos[1] - cx
    d = np.sqrt(dx**2 + dy**2)

    freq_px = d / max(h, w)
    spacing_px = 1 / freq_px
    spacing_nm = spacing_px * nm_per_pixel

    print(f"🧠 [FFT] 主频距中心：{d:.2f} px → 间距：{spacing_nm:.2f} nm")
    return spacing_nm


def visualize_results(image, profiles, peaks_list, avg_spacing_nm, save_path):
    fig, ax = plt.subplots(1, 2, figsize=(12, 5))
    ax[0].imshow(image, cmap='gray')
    ax[0].set_title("Input Image")
    ax[0].axis('off')

    for i, profile in enumerate(profiles):
        ax[1].plot(255 - profile, label=f"Line {i+1}")
        if i < len(peaks_list):
            peaks = peaks_list[i]
            ax[1].plot(peaks, (255 - profile)[peaks], "x")

    ax[1].set_title(f"Avg spacing ≈ {avg_spacing_nm:.2f} nm")
    ax[1].set_xlabel("Pixel")
    ax[1].set_ylabel("Intensity")
    ax[1].legend()
    ax[1].grid(True)

    plt.tight_layout()
    plt.savefig(save_path, dpi=300)
    plt.show()


def main(image_path):
    img = cv2.imread(image_path)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    nm_per_pixel = detect_scale_bar(gray)
    best_freq, filtered = find_best_frequency(gray, nm_per_pixel)

    avg_nm, spacings, profiles, peaks_list = extract_profile_and_spacing(filtered, nm_per_pixel)

    print("\n📊 结果对比：")
    if avg_nm:
        print(f"🔹 Gabor条纹间距：{avg_nm:.2f} nm")
    else:
        print(f"🔹 Gabor未能检测到足够条纹")

    spacing_fft = extract_spacing_fft(gray, nm_per_pixel)
    print(f"🔸 FFT条纹间距：{spacing_fft:.2f} nm")

    output_path = os.path.join("../outputs", os.path.splitext(os.path.basename(image_path))[0] + "_result.png")
    os.makedirs("../outputs", exist_ok=True)
    visualize_results(gray, profiles, peaks_list, avg_nm or 0.0, output_path)


if __name__ == "__main__":
    image_path = "../img/pic3.jpg"  # 修改为你的图片路径
    main(image_path)